import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter
from scipy.ndimage import gaussian_filter, uniform_filter

def ssim(img1, img2):
    """
    计算两幅灰度图像的SSIM（结构相似性度量）。

    参数:
    img1, img2 (numpy arrays): 输入的两幅灰度图像，数据类型应为float且值域在[0, 1]之间。

    返回:
    float: SSIM值，范围在[-1, 1]之间，值越大表示图像越相似。
    """
    # 输入检查
    if img1.max() > 1 or img2.max() > 1:
        raise ValueError("Inputs should be normalized to the range [0, 1].")

    C1 = (0.01 * 255) ** 2
    C2 = (0.03 * 255) ** 2

    #当前公式已经有 C1 和 C2，它们本身可以防止分母为零
    #eps = 1e-10  # 添加一个小的正数以防止除以零

    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    
    #用不到这3个值
    # K1 = 0.01
    # K2 = 0.03
    # L = 255  # 灰度级别，对于8位灰度图像

    # 选择滤波器,让函数适应更多场景
    filter_func = gaussian_filter if use_gaussian else uniform_filter

    # 计算均值，window_size (int): 滤波器窗口的大小，默认为11
    mu1 = filter_func(img1, window_size)
    mu2 = filter_func(img2, window_size)

    # 计算方差和协方差
    mu1_sq = mu1 * mu1
    mu2_sq = mu2 * mu2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = filter_func(img1 ** 2, window_size) - mu1_sq
    sigma2_sq = filter_func(img2 ** 2, window_size) - mu2_sq
    sigma12 = filter_func(img1 * img2, window_size) - mu1_mu2

    # 计算SSIM
    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / (
        (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
    )

    # 返回整个图像的平均值作为SSIM值
    return np.mean(ssim_map)


def preprocess_image(img_path):
    """
    加载图像，转换为灰度图，并归一化到[0, 1]范围。
    """
    img = Image.open(img_path).convert("L")  # 转换为灰度图
    img = np.array(img, dtype=np.float64) / 255.0  # 归一化到[0, 1]范围
    return img


# 示例：使用SSIM函数
img1_path = "C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\5.png"  # 替换为你的第一张图像路径
img2_path = "C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\ans.png"  # 替换为你的第二张图像路径

img1 = preprocess_image(img1_path)
img2 = preprocess_image(img2_path)

# 确保两张图像尺寸相同
assert img1.shape == img2.shape, "Images must have the same dimensions."

# 计算SSIM值
ssim_value = ssim(img1, img2)

# 打印SSIM值
print(f"The SSIM value between the two images is: {ssim_value}")
